﻿#include "FrameRenderer.h"
#include <QSGRendererInterface>
#include <QQuickWindow>
#include <QSize>
#include <QRunnable>

GLRenderer::GLRenderer(QObject *parent)
    :QObject(parent)

{
    
}

void GLRenderer::clearRenderContext()
{
    av_frame_free(&m_curFrame);

    if(m_yTex)
        glDeleteTextures(1, &m_yTex);

    if(m_vTex)
        glDeleteTextures(1, &m_vTex);

    if(m_uTex)
        glDeleteTextures(1, &m_uTex);
}

GLRenderer::~GLRenderer()
{
    clearRenderContext();

    glDeleteBuffers(1, &m_vbo);
    delete m_shaderProgram;
}

void GLRenderer::setFrame(AVFrame * frame)
{
    if(!frame)
        return;

    if(m_nextFrame)
        return;
        // av_frame_free(&m_nextFrame);

    m_nextFrame = av_frame_clone(frame);
}

void GLRenderer::init()
{
    QSGRendererInterface *rif = window()->rendererInterface();
    Q_ASSERT(rif->graphicsApi() == QSGRendererInterface::OpenGL);

    initializeGL();
}

void GLRenderer::makeYUVTextures(int yl, int ul, int vl, int h)
{
    // 创建Y分量纹理
    glGenTextures(1, &m_yTex);
    glBindTexture(GL_TEXTURE_2D, m_yTex);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_LUMINANCE, yl, h, 0, GL_LUMINANCE, GL_UNSIGNED_BYTE, nullptr);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glBindTexture(GL_TEXTURE_2D, 0);

    // 创建U分量纹理，尺寸为Y分量的1/4
    glGenTextures(1, &m_uTex);
    glBindTexture(GL_TEXTURE_2D, m_uTex);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_LUMINANCE, ul, h / 2, 0, GL_LUMINANCE, GL_UNSIGNED_BYTE, nullptr);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glBindTexture(GL_TEXTURE_2D, 0);

    // 创建V分量纹理，尺寸为Y分量的1/4
    glGenTextures(1, &m_vTex);
    glBindTexture(GL_TEXTURE_2D, m_vTex);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_LUMINANCE, vl, h / 2, 0, GL_LUMINANCE, GL_UNSIGNED_BYTE, nullptr);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glBindTexture(GL_TEXTURE_2D, 0);
}

void GLRenderer::initializeGL()
{
    if(m_shaderProgram)
        return;


    QSGRendererInterface *rif = m_window->rendererInterface();
    Q_ASSERT(rif->graphicsApi() == QSGRendererInterface::OpenGL);
    initializeOpenGLFunctions();

    // 顶点着色器源码
    const char* vertexShaderSource = R"(attribute vec4 position;
        attribute vec2 yTexCoord;
        attribute vec2 uTexCoord;
        attribute vec2 vTexCoord;
        varying vec2 v_yTexCoord;
        varying vec2 v_uTexCoord;
        varying vec2 v_vTexCoord;

        void main() {
            gl_Position = position;
            v_yTexCoord = yTexCoord;
            v_uTexCoord = uTexCoord;
            v_vTexCoord = vTexCoord;
        }
    )";

    // 片元着色器源码
    const char* fragmentShaderSource = R"(
        varying vec2 v_yTexCoord;
        varying vec2 v_uTexCoord;
        varying vec2 v_vTexCoord;
        uniform sampler2D yTex;
        uniform sampler2D uTex;
        uniform sampler2D vTex;

        void main() {
            vec3 yuv, rgb;
            vec3 yuv2r = vec3(1.164, 0.0, 1.596);
            vec3 yuv2g = vec3(1.164, -0.391, -0.813);
            vec3 yuv2b = vec3(1.164, 2.018, 0.0);

            yuv.x = texture2D(yTex, v_yTexCoord).r - 0.0625;
            yuv.y = texture2D(uTex, v_uTexCoord).r - 0.5;
            yuv.z = texture2D(vTex, v_vTexCoord).r - 0.5;

            rgb.x = dot(yuv, yuv2r);
            rgb.y = dot(yuv, yuv2g);
            rgb.z = dot(yuv, yuv2b);

            gl_FragColor = vec4(rgb, 1.0);
        }
    )";

    // 创建着色器程序
    m_shaderProgram = new QOpenGLShaderProgram();

    // 编译顶点着色器
    QOpenGLShader* vertexShader = new QOpenGLShader(QOpenGLShader::Vertex);
    if (!vertexShader->compileSourceCode(vertexShaderSource)) {
        qWarning() << "Failed to compile vertex shader: " << vertexShader->log();
        delete vertexShader;
        return;
    }

    // 编译片元着色器
    QOpenGLShader* fragmentShader = new QOpenGLShader(QOpenGLShader::Fragment);
    if (!fragmentShader->compileSourceCode(fragmentShaderSource)) {
        qWarning() << "Failed to compile fragment shader: " << fragmentShader->log();
        delete fragmentShader;
        delete vertexShader;
        return;
    }

    // 链接着色器程序
    m_shaderProgram->addShader(vertexShader);
    m_shaderProgram->addShader(fragmentShader);
    if (!m_shaderProgram->link()) {
        qWarning() << "Failed to link shader program: " << m_shaderProgram->log();
        delete m_shaderProgram;
        m_shaderProgram = nullptr;
        delete fragmentShader;
        delete vertexShader;
        return;
    }

    // 创建VBO
    glGenBuffers(1, &m_vbo);
}

void GLRenderer::setWindow(QQuickWindow *window)
{
    m_window = window;
}

void GLRenderer::resizeGL(int w, int h)
{
    m_viewportWidth = w;
    m_viewportHeight = h;
}

void GLRenderer::paintFrame()
{
    window()->beginExternalCommands();

    if(!m_curFrame && !m_nextFrame) {
        window()->endExternalCommands();
        return;
    }

    if(m_nextFrame){
        deliverFrame(m_nextFrame);
        av_frame_free(&m_curFrame);
        m_curFrame = m_nextFrame;
        m_nextFrame = nullptr;
    }

    if(!m_curFrame) {
        window()->endExternalCommands();
        return;
    }

    glViewport(0, 0, m_viewportWidth, m_viewportHeight);

    // 绑定并使用着色器程序
    m_shaderProgram->bind();

    // 绑定VBO
    glBindBuffer(GL_ARRAY_BUFFER, m_vbo);


    // 在保持长宽比下计算图像帧的顶点位置
    float top = 1.0f;
    float bottom = -1.0f;
    float left = -1.0f;
    float right = 1.0f;
    QSizeF frameSize (m_curFrame->width, m_curFrame->height);
    auto viewportSize = QSizeF(m_viewportWidth, m_viewportHeight);
    frameSize.scale(viewportSize, Qt::AspectRatioMode::KeepAspectRatio);

    if(frameSize.width() == viewportSize.width()) {
        top = top - (2.0f - float(frameSize.height()) / viewportSize.height() * 2.0f) / 2.0f;
        bottom = bottom + (2.0f - float(frameSize.height()) / viewportSize.height() * 2.0f) / 2.0f;
    }else { // if(m_curFrame->height == 480)
        left = left + (2.0f - float(frameSize.width()) / viewportSize.width() * 2.0f) / 2.0f;
        right = right - (2.0f - float(frameSize.width()) / viewportSize.width() * 2.0f) / 2.0f;
    }

    // 滤除纹理中因为内存对齐而导致的纹理右边的空白
    float yTexcoordsU = float(m_curFrame->width) / m_curFrame->linesize[0];
    float uTexcoordsU = float(m_curFrame->width) / 2 / m_curFrame->linesize[1];
    float vTexcoordsU = float(m_curFrame->width) / 2 / m_curFrame->linesize[2];

    float vertices[] = {
        // 位置                    Y纹理坐标             U纹理坐标            V纹理坐标
        left, bottom, 0.0f,       0.0f, 1.0f,           0.0f, 1.0f,         0.0f, 1.0f,  // 左下
        right, bottom, 0.0f,      yTexcoordsU, 1.0f,    uTexcoordsU, 1.0f,  vTexcoordsU, 1.0f,// 右下
        right,  top, 0.0f,        yTexcoordsU, 0.0f,    uTexcoordsU, 0.0f,  vTexcoordsU, 0.0f,// 右上
        right,  top, 0.0f,        yTexcoordsU, 0.0f,    uTexcoordsU, 0.0f,  vTexcoordsU, 0.0f,// 右上
        left,  top, 0.0f,         0.0f, 0.0f,           0.0f, 0.0f,         0.0f, 0.0f,// 左上
        left, bottom, 0.0f,       0.0f, 1.0f,           0.0f, 1.0f,         0.0f, 1.0f,// 左下
    };
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

    // 获取着色器中的位置和纹理坐标属性位置
    GLint posAttr = m_shaderProgram->attributeLocation("position");
    GLint yTexCoordAttr = m_shaderProgram->attributeLocation("yTexCoord");
    GLint uTexCoordAttr = m_shaderProgram->attributeLocation("uTexCoord");
    GLint vTexCoordAttr = m_shaderProgram->attributeLocation("vTexCoord");

    // 设置顶点位置属性
    glEnableVertexAttribArray(posAttr);
    glVertexAttribPointer(posAttr, 3, GL_FLOAT, GL_FALSE, 9 * sizeof(float), (void*)0);

    // 设置纹理坐标属性
    glEnableVertexAttribArray(yTexCoordAttr);
    glVertexAttribPointer(yTexCoordAttr, 2, GL_FLOAT, GL_FALSE, 9 * sizeof(float), (void*)(3 * sizeof(float)));

    glEnableVertexAttribArray(uTexCoordAttr);
    glVertexAttribPointer(uTexCoordAttr, 2, GL_FLOAT, GL_FALSE, 9 * sizeof(float), (void*)(5 * sizeof(float)));

    glEnableVertexAttribArray(vTexCoordAttr);
    glVertexAttribPointer(vTexCoordAttr, 2, GL_FLOAT, GL_FALSE, 9 * sizeof(float), (void*)(7 * sizeof(float)));

    // 绑定并激活YUV纹理
    glActiveTexture(GL_TEXTURE0);
    glBindTexture(GL_TEXTURE_2D, m_yTex);
    m_shaderProgram->setUniformValue("yTex", 0);

    glActiveTexture(GL_TEXTURE1);
    glBindTexture(GL_TEXTURE_2D, m_uTex);
    m_shaderProgram->setUniformValue("uTex", 1);

    glActiveTexture(GL_TEXTURE2);
    glBindTexture(GL_TEXTURE_2D, m_vTex);
    m_shaderProgram->setUniformValue("vTex", 2);

    glDisable(GL_DEPTH_TEST);

    glEnable(GL_BLEND);
    
    glBlendFunc(GL_SRC_ALPHA, GL_ONE);

    // 绘制矩形
    glDrawArrays(GL_TRIANGLES, 0, 6);

    // 禁用顶点属性
    glDisableVertexAttribArray(posAttr);
    glDisableVertexAttribArray(yTexCoordAttr);
    glDisableVertexAttribArray(uTexCoordAttr);
    glDisableVertexAttribArray(vTexCoordAttr);

    // 解绑VBO
    glBindBuffer(GL_ARRAY_BUFFER, 0);

    // 解绑着色器程序
    m_shaderProgram->release();

    window()->endExternalCommands();
}

/**
 * @brief 将当前的AVFrame更新到GL纹理
 */
void GLRenderer::deliverFrame(AVFrame* frame)
{
    int yl = frame->linesize[0];
    int ul = frame->linesize[1];
    int vl = frame->linesize[2];
    int height = frame->height;

    if(m_yTex == 0)
        makeYUVTextures(yl, ul, vl, height);

    // 更新Y分量纹理
    glBindTexture(GL_TEXTURE_2D, m_yTex);
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, yl, height, GL_LUMINANCE, GL_UNSIGNED_BYTE, frame->data[0]);
    glFinish();

    // 更新U分量纹理，尺寸为Y分量的一半
    glBindTexture(GL_TEXTURE_2D, m_uTex);
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, ul, height / 2, GL_LUMINANCE, GL_UNSIGNED_BYTE, frame->data[1]);
    glFinish();

    // 更新V分量纹理，尺寸为Y分量的一半
    glBindTexture(GL_TEXTURE_2D, m_vTex);
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, vl, height / 2, GL_LUMINANCE, GL_UNSIGNED_BYTE, frame->data[2]);
    glFinish();

    // 解绑纹理
    glBindTexture(GL_TEXTURE_2D, 0);
}

FrameRenderer::FrameRenderer()
{
    connect(this, &QQuickItem::windowChanged, this, &FrameRenderer::handleWindowChanged);
}

void FrameRenderer::handleWindowChanged(QQuickWindow *win)
{
    if (win) {
        connect(win, &QQuickWindow::beforeSynchronizing, this, &FrameRenderer::sync, Qt::DirectConnection);
        connect(win, &QQuickWindow::sceneGraphInvalidated, this, &FrameRenderer::cleanup, Qt::DirectConnection);

        // Ensure we start with cleared to black. The squircle's blend mode relies on this.
        win->setColor(Qt::black);
    }
}

class CleanupJob : public QRunnable
{
public:
    CleanupJob(GLRenderer *renderer) : m_renderer(renderer) { }
    void run() override { delete m_renderer; }
private:
    GLRenderer *m_renderer;
};

void FrameRenderer::releaseResources()
{
    window()->scheduleRenderJob(new CleanupJob(m_glRenderer), QQuickWindow::BeforeSynchronizingStage);
    m_glRenderer = nullptr;
}

void FrameRenderer::sync()
{
    if(!m_glRenderer) {
        m_glRenderer = new GLRenderer();
        connect(window(), &QQuickWindow::beforeRendering, m_glRenderer, &GLRenderer::init, Qt::DirectConnection);
        connect(window(), &QQuickWindow::beforeRenderPassRecording, m_glRenderer, &GLRenderer::paintFrame, Qt::DirectConnection);
    }

    m_glRenderer->setWindow(window());

    auto viewportSize = window()->size() * window()->devicePixelRatio();
    m_glRenderer->resizeGL(viewportSize.width(), viewportSize.height());

    window()->setColor(Qt::black);
}

void FrameRenderer::cleanup()
{
    if(m_glRenderer) {
        delete m_glRenderer;
        m_glRenderer = nullptr;
    }
}